/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_jacobi_preconditioner.h"

namespace YHAMG
{

void ParCSR_Jacobi_Preconditioner::Free()
{
	D_Recip.Free();
}

void ParCSR_Jacobi_Preconditioner::Setup(const ParCSR_Matrix& A, int REUSE)
{
	int n = A.LocLoc.size[0];
	int* Ap = A.LocLoc.rowptr;
	int* Ai = A.LocLoc.colind;
	double* Av = A.LocLoc.values;

	if (REUSE == 0)
	{
		comm = A.comm;
		D_Recip.Resize(n);
	}

	double* Dv_Recip = D_Recip.values;

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		for (int j = Ap[i]; j < Ap[i + 1]; ++j)
		{
			if (Ai[j] == i)
			{
				Dv_Recip[i] = 1.0 / Av[j];
				break;
			}
		}
	}
}

int ParCSR_Jacobi_Preconditioner::InSize() const
{
	return D_Recip.size;
}

int ParCSR_Jacobi_Preconditioner::OutSize() const
{
	return D_Recip.size;
}

void ParCSR_Jacobi_Preconditioner::Apply(const Par_Vector& b, const Par_Vector& x) const
{
	int n = D_Recip.size;
	double* Dv_Recip = D_Recip.values;
	double* xv = x.Loc.values;
	double* bv = b.Loc.values;

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
	for (int i = 0; i < n; ++i)
		xv[i] = bv[i] * Dv_Recip[i];
}

}